package com.kara.woodAgent.agent.wood;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.kara.woodAgent.agent.graph.BaseGraph;
import com.kara.woodAgent.agent.graph.Next;
import com.kara.woodAgent.agent.graph.impl.SubGraphNodeImpl;
import com.kara.woodAgent.agent.model.ModelProvider;
import com.kara.woodAgent.agent.tool.provider.ToolProvider;
import com.kara.woodAgent.agent.wood.context.WoodContext;
import com.kara.woodAgent.agent.wood.context.WoodReactContext;
import com.kara.woodAgent.agent.wood.model.Plan;
import com.kara.woodAgent.agent.wood.model.SpecialAgentResult;
import com.kara.woodAgent.agent.wood.node.ActionToolNode;
import com.kara.woodAgent.agent.wood.node.SummaryNode;
import com.kara.woodAgent.agent.wood.node.ThinkNode;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;

import java.util.ArrayList;
import java.util.List;

public class SpecialAgentTemplate {


    public static SubGraphNodeImpl<WoodContext, WoodReactContext> createSubGraphNode(String consumePrompt, ModelProvider modelProvider, ToolProvider.ToolServiceHolder toolServiceHolder) {
        //正常情况下 lastMessage是 退出原因

        BaseGraph<WoodReactContext> reactGraph = new BaseGraph<WoodReactContext>()
                .addStartNode("model", new ThinkNode(modelProvider, toolServiceHolder))
                .addNode("tool", new ActionToolNode(toolServiceHolder))
                .addNode("summary", new SummaryNode(modelProvider));

        return new SubGraphNodeImpl<WoodContext, WoodReactContext>(reactGraph)
                .gcSupplier(woodContext -> {

                    Plan currentPlan = woodContext.getCurrentPlan();

                    String subQuestion = currentPlan.getContent();

                    String systemPrompt = """
                            你是wood智能体系统的的子节点之一, 角色是：专项智能体。
                            #### 职责
                            你的职责是 结合分配给你的任务,完成分配给你的任务
                            来完成任务.
                            {{consumePrompt}}
                            #### 你的任务是 ####
                            {{question}}
                            """;
                    systemPrompt = systemPrompt.replace("{{consumePrompt}}", consumePrompt);
                    List<String> knowList = woodContext.getKnowList();
                    if (!knowList.isEmpty()) {
                        //添加已知信息
                        systemPrompt = systemPrompt + "\n#### 已知信息：####\n" + String.join("\n", knowList);
                    }

                    SystemMessage systemMessage = SystemMessage.from(systemPrompt.replace("{{question}}", subQuestion));

                    WoodReactContext reactGraphContext = new WoodReactContext();
                    reactGraphContext.setPlan(subQuestion);
                    List<ChatMessage> messages = new ArrayList<>();
                    messages.add(systemMessage);
                    reactGraphContext.setChatMessages(messages);

                    return reactGraphContext;
                })
                .next((woodContext, reactGraphContext) -> {

                    //正常情况下 lastMessage是 退出原因
                    SpecialAgentResult result = reactGraphContext.getSummary();

                    Plan currentPlan = woodContext.getCurrentPlan();
                    currentPlan.setStatus(result.status());
                    if (Plan.SUCCESS.equals(currentPlan.getStatus())) {
                        currentPlan.setResultMsg(result.answer());
                    }
                    currentPlan.setMessage(result.message());

                    return Next.NextNode("dispatch").step(JSON.toJSONString(currentPlan));
                });
    }


    public static SubGraphNodeImpl<WoodContext, WoodReactContext> createSubGraphNode(ModelProvider modelProvider, ToolProvider.ToolServiceHolder toolServiceHolder) {
        return createSubGraphNode("", modelProvider, toolServiceHolder);

    }
}
